#! /home/coder/miniconda3/bin/python
import re
import os
import json

def write_json(path,info):
    with open(path,"w",encoding='utf8') as f:
        f.write(json.dumps(info, ensure_ascii=False,indent=4, separators=(',', ':')))
    return 'ok'

def read_json(path):
    with open(path,'r',encoding='utf8')as fp:
        json_data = json.load(fp)
    return json_data


# 建立 bed 文件索引，保存在 bedindex.json 中
def builtBedInedx(file):
    bedindex = dict()
    with open(file,"r",encoding="utf8") as f:
        line_num = 1
        line = f.readline().strip().split('\t')
        chr_before = re.findall(f"([cC]hr)?(.*)$",line[0])[0][1]
        strand_before = line[4]
        bedindex[chr_before]=dict()
        bedindex[chr_before][line[4]]=[line_num,None]

        while True:
            line = f.readline().strip().split('\t')
            line_num += 1

            if line == [""]:
                bedindex[chr_before][strand_before][1]=line_num-1
                bedindex["total"]=line_num-1
                break

            chr_num = re.findall(f"([cC]hr)?(.*)$",line[0])[0][1]

            if chr_num not in bedindex:
                bedindex[chr_num]=dict()

            if line[4] not in bedindex[chr_num]:
                bedindex[chr_num][line[4]]=[line_num,None]
                bedindex[chr_before][strand_before][1]=line_num-1

            chr_before = chr_num
            strand_before = line[4]
            
    write_json("bedindex.json",bedindex)


# 二分查找法
def getGeneBed(file,tofile,chr="1",inbed=1,outbed=1,minOverlap=1,strand="-"):
    bedindex = read_json("bedindex.json")
    total=bedindex["total"]
    left,right=bedindex[str(chr)][strand]
    inbed+=minOverlap-1
    outbed-=minOverlap+1

    def search(a,b):
        if b-a<=1:
            return (a,b)[mode]
        x = _getTail(a)
        y = _getTail(b)
        mid = (a+b)//2
        z = _getTail(mid)
        if target>z:
            return search(mid,b)
        else:
            return search(a,mid)
    
    def _getTail(line_num):
        line = os.popen('sed -n {}p {}'.format(line_num, file)).read().strip().split('\t')
        return int(line[mode+1])

    result=list()
    for target,mode in zip((inbed,outbed),(1,0)):
        result.append(search(left,right))
    
    a,b=result
    if a>b:
        with open(tofile,'w',encoding="utf-8",newline="") as f:
            f.write(f"# geneX: 0; total: {total}\n")
    else:
        with open(tofile,'w',encoding="utf-8",newline="") as f:
            data = os.popen('''sed -n {},{}p {} | awk -F "\t" '$3-$2 >= {}'
                                '''.format(a,b,file,minOverlap)).read()
            data = data.strip()
            lines = len(data.split("\n")) if data else 0
            f.write('# geneX: {}; total: {}\n'.format(lines,total))
            f.write(data)
        
    return "Finish! Please cheak on {}".format(tofile)


if __name__ == "__main__":
    builtBedInedx("chr.bed")
    print(getGeneBed("chr.bed","geneX_chr3.bed",chr="3"
               ,inbed=20034100,outbed=20147100))
    print(getGeneBed("chr.bed","geneX_chr3_100.bed",chr="3"
               ,inbed=20034100,outbed=20147100,minOverlap=100))
    print(getGeneBed("chr.bed","geneX_chr3_1000.bed",chr="3"
               ,inbed=20034100,outbed=20147100,minOverlap=1000))
    print(getGeneBed("chr.bed","geneX_chr5.bed",chr="5"
               ,inbed=20034100,outbed=20147100,strand="+"))